#include "ShaderReflector.h"
#include <winnt.h>
#include <string>
#include <iostream>
#include "dxcapi.h"
namespace SL
{

ShaderReflector::ShaderReflector()
{
    // Create compiler and utils.
    //
    DxcCreateInstance(CLSID_DxcUtils, IID_PPV_ARGS(&pUtils));
    DxcCreateInstance(CLSID_DxcCompiler, IID_PPV_ARGS(&pCompiler));

    // Create default include handler. (You can create your own...)
    //
    pUtils->CreateDefaultIncludeHandler(&pIncludeHandler);
}

ShaderReflector::~ShaderReflector()
{
    if (pIncludeHandler)
    {
        pIncludeHandler->Release();
        pIncludeHandler = nullptr;
    }
    if (pCompiler)
    {
        pCompiler->Release();
        pCompiler = nullptr;
    }
    if (pUtils)
    {
        pUtils->Release();
        pUtils = nullptr;
    }
}

std::vector<LPCWSTR> ShaderReflector::buildCompileArgs(LPCWSTR shaderFile, ShaderType shaderType, bool debug = true)
{
    std::vector<LPCWSTR> args;
    std::wstring suffix;
    switch (shaderType)
    {
    case ShaderType::Vertex:
        args.push_back(L"-E");
        args.push_back(L"VSMain");
        args.push_back(L"-T");
        args.push_back(L"vs_6_0"); // Target for vertex shader.
        suffix = L".vs";
        break;
    case ShaderType::Pixel:
        args.push_back(L"-E");
        args.push_back(L"PSMain");
        args.push_back(L"-T");
        args.push_back(L"ps_6_0"); // Target for pixel shader.
        suffix = L".ps";
        break;
    case ShaderType::Geometry:
        // TODO
        break;
    case ShaderType::Compute:
        args.push_back(L"-E");
        args.push_back(L"CSMain");
        args.push_back(L"-T");
        args.push_back(L"cs_6_0"); // Target for compute shader.
        suffix = L".cs";
        break;
    case ShaderType::RayTracing:
        args.push_back(L"-T");
        args.push_back(L"lib_6_6"); // Target for raytracing shader library.
        suffix = L".rt";
        break;
    default:
        break;
    }
    if (debug)
    {
        args.push_back(L"-Zi"); // Enable debug information
        // args.push_back(L"-Zs");  // Enable Debug (slim format)
        args.push_back(L"-Fd");
        args.push_back((std::wstring(shaderFile) + suffix + L".pdb").c_str());
        args.push_back(L"-Fo");
        args.push_back((std::wstring(shaderFile) + suffix + L".cso").c_str());
        args.push_back(L"-Qstrip_debug");
    }
    args.push_back(L"-Qstrip_reflect"); // Strip reflection into a separate blob.
    args.push_back(DXC_ARG_ALL_RESOURCES_BOUND);


    return args;
}

ID3DBlob* ShaderReflector::compileShaderToObj(LPCWSTR shaderFile, ShaderType shaderType = ShaderType::Unknown)
{
    CComPtr<IDxcResult> pResults = compileShader(shaderFile, shaderType);
    IDxcBlob* pObj;
    pResults->GetOutput(DXC_OUT_OBJECT, IID_PPV_ARGS(&pObj), nullptr);
    return reinterpret_cast<ID3DBlob*>(pObj); // just type aliasing
}

IDxcBlobUtf8* ShaderReflector::compileShaderToDisasm(LPCWSTR shaderFile, ShaderType shaderType = ShaderType::Unknown)
{
    CComPtr<IDxcResult> pResults = compileShader(shaderFile, shaderType);
    IDxcBlobUtf8* pDisasm;
    pResults->GetOutput(DXC_OUT_DISASSEMBLY, IID_PPV_ARGS(&pDisasm), nullptr);
    return pDisasm; // just type aliasing
}

CComPtr<IDxcResult> ShaderReflector::compileShader(LPCWSTR shaderFile, ShaderType shaderType = ShaderType::Unknown)
{
    // Open source file.
    //
    CComPtr<IDxcBlobEncoding> pSource = nullptr;
    HRESULT hr = pUtils->LoadFile(shaderFile, nullptr, &pSource);
    if (hr != S_OK)
    {
        wprintf(L"Failed to load shader file: %s\n", shaderFile);
        return {};
    }

    DxcBuffer Source = {
        .Ptr = pSource->GetBufferPointer(),
        .Size = pSource->GetBufferSize(),
        .Encoding = DXC_CP_ACP // Assume BOM says UTF8 or UTF16 or this is ANSI text.
    };

    if (shaderType == ShaderType::Unknown)
        shaderType = getShaderType(shaderFile);  // Function to determine shader type based on file extension or content.
    auto pszArgs = buildCompileArgs(shaderFile, shaderType); // Change to desired shader type.

    // Compile it with specified arguments.
    //
    CComPtr<IDxcResult> pResults;
    pCompiler->Compile(&Source, pszArgs.data(), pszArgs.size(), pIncludeHandler, IID_PPV_ARGS(&pResults));

    // Print errors if present.
    //
    CComPtr<IDxcBlobUtf8> pErrors = nullptr;
    pResults->GetOutput(DXC_OUT_ERRORS, IID_PPV_ARGS(&pErrors), nullptr);
    if (pErrors != nullptr && pErrors->GetStringLength() != 0)
        wprintf(L"Warnings and Errors:\n%S\n", pErrors->GetStringPointer());

    return pResults;
}

void ShaderReflector::reflectShader(LPCWSTR shaderFile)
{
    CComPtr<IDxcResult> pResults = compileShader(shaderFile);

    // Quit if the compilation failed.
    //
    HRESULT hrStatus;
    pResults->GetStatus(&hrStatus);
    if (FAILED(hrStatus))
    {
        wprintf(L"Compilation Failed\n");
        return;
    }

    CComPtr<IDxcBlob> pReflectionData;
    pResults->GetOutput(DXC_OUT_REFLECTION, IID_PPV_ARGS(&pReflectionData), nullptr);

    std::wcout << shaderFile << std::endl;

    ShaderType shaderType = getShaderType(shaderFile); // Function to determine shader type based on file extension or content.
    printResourceBindings(shaderType, pReflectionData, pUtils);
}

void ShaderReflector::printResourceBindings(ShaderType shaderType, IDxcBlob* pReflectionData, IDxcUtils* pUtils)
{
    DxcBuffer ReflectionData = {
        .Ptr = pReflectionData->GetBufferPointer(),
        .Size = pReflectionData->GetBufferSize(),
        .Encoding = DXC_CP_ACP, // for binary
    };
    if (shaderType == ShaderType::RayTracing)
    {
        // raytracing
        CComPtr<ID3D12LibraryReflection> pReflection;
        pUtils->CreateReflection(&ReflectionData, IID_PPV_ARGS(&pReflection));

        if (pReflection != nullptr)
        {
            D3D12_LIBRARY_DESC libdesc;
            pReflection->GetDesc(&libdesc);

            for (auto i = 0; i < libdesc.FunctionCount; i++)
            {
                ID3D12FunctionReflection* functionReflect = pReflection->GetFunctionByIndex(i);
                D3D12_FUNCTION_DESC shader_desc;
                functionReflect->GetDesc(&shader_desc);

                for (int j = 0; j < shader_desc.BoundResources; j++)
                {
                    D3D12_SHADER_INPUT_BIND_DESC resource_desc;
                    functionReflect->GetResourceBindingDesc(j, &resource_desc);

                    LPCSTR shaderVarName = resource_desc.Name;
                    UINT registerSpace = resource_desc.Space;
                    D3D_SHADER_INPUT_TYPE resourceType = resource_desc.Type;
                    UINT bindPoint = resource_desc.BindPoint;

                    std::cout << to_string(resourceType) << " : " << shaderVarName << " (bind point: " << bindPoint
                              << ", register space: " << registerSpace << ")" << std::endl;
                    std::cout << "------------------------" << std::endl;
                }
            }
        }
    }
    else
    {
        // non raytracing
        CComPtr<ID3D12ShaderReflection> shaderReflection;
        pUtils->CreateReflection(&ReflectionData, IID_PPV_ARGS(&shaderReflection));

        if (shaderReflection != nullptr)
        {
            D3D12_SHADER_DESC shaderDesc;
            shaderReflection->GetDesc(&shaderDesc);

            for (uint32_t i = 0; i < shaderDesc.BoundResources; i++)
            {
                D3D12_SHADER_INPUT_BIND_DESC resourceDesc;
                shaderReflection->GetResourceBindingDesc(i, &resourceDesc);

                auto shaderVarName = resourceDesc.Name;
                auto registerSpace = resourceDesc.Space;
                auto resourceType = resourceDesc.Type;
                auto bindPoint = resourceDesc.BindPoint;

                std::cout << to_string(resourceType) << " : " << shaderVarName << " (bind point: " << bindPoint
                          << ", register space: " << registerSpace << ")" << std::endl;
                std::cout << "------------------------" << std::endl;
            }
        }
    }
}

ShaderReflector::ShaderType ShaderReflector::getShaderType(LPCWSTR shaderFile)
{
    // Determine shader type based on file extension or content.
    // For simplicity, we assume the file name contains the type.
    std::wstring fileName(shaderFile);
    if (fileName.find(L".vert") != std::wstring::npos)
        return ShaderType::Vertex;
    else if (fileName.find(L".frag") != std::wstring::npos)
        return ShaderType::Pixel;
    else if (fileName.find(L".comp") != std::wstring::npos)
        return ShaderType::Compute;
    else if (fileName.find(L".rayt") != std::wstring::npos)
        return ShaderType::RayTracing;
    else if (fileName.find(L".geom") != std::wstring::npos)
        return ShaderType::Geometry;
    // Add more cases as needed.
    return ShaderType::Unknown; // Default case.
}
} // namespace SL